package org.wenshu.ai.modular.chat.provider.core.oceanbase;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;

import java.sql.*;
import java.util.*;

public class OceanbaseEmbeddingStore implements EmbeddingStore<TextSegment> {

    private final String url;
    private final String username;
    private final String password;
    private final String tableName;
    private final int dimension;

    public OceanbaseEmbeddingStore(String url, String username, String password, String tableName, int dimension) {
        this.url = url;
        this.username = username;
        this.password = password;
        this.tableName = tableName;
        this.dimension = dimension;
        initializeTable();
    }

    private void initializeTable() {
        try (Connection conn = getConnection()) {
            String createTableSql = String.format(
                "CREATE TABLE IF NOT EXISTS %s (" +
                "id VARCHAR(36) PRIMARY KEY, " +
                "text TEXT, " +
                "metadata TEXT, " +
                "embedding VECTOR(%d), " +
                "VECTOR INDEX idx_embedding(embedding) WITH (distance=L2, type=hnsw)" +
                ")", tableName, dimension);

            try (Statement stmt = conn.createStatement()) {
                stmt.execute(createTableSql);
            }
        } catch (SQLException e) {
            throw new RuntimeException("Failed to initialize table", e);
        }
    }

    @Override
    public String add(Embedding embedding) {
        List<String> ids = addAll(List.of(embedding));
        return ids.get(0);
    }

    @Override
    public void add(String id, Embedding embedding) {
        addAll(List.of(id), List.of(embedding), null);
    }

    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        List<String> ids = addAll(List.of(embedding), List.of(textSegment));
        return ids.get(0);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        return addAll(embeddings, null);
    }

    @Override
    public void addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> embedded) {
        if (ids.size() != embeddings.size() || (embedded != null && embedded.size() != embeddings.size())) {
            throw new IllegalArgumentException("Size of ids, embeddings and embedded must match");
        }
        try (Connection conn = getConnection()) {
            String sql = String.format(
                "REPLACE INTO %s (id, text, metadata, embedding) VALUES (?, ?, ?, ?)",
                tableName
            );

            try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
                Gson gson = new Gson();
                for (int i = 0; i < embeddings.size(); i++) {
                    TextSegment textSegment = embedded != null ? embedded.get(i) : null;
                    
                    pstmt.setString(1, ids.get(i));
                    pstmt.setString(2, textSegment != null ? textSegment.text() : null);
                    pstmt.setString(3, textSegment != null && textSegment.metadata() != null ?
                        gson.toJson(textSegment.metadata().asMap()) : "{}");
                    pstmt.setString(4, formatVector(embeddings.get(i).vector()));
                    pstmt.addBatch();
                }
                pstmt.executeBatch();
            }
        } catch (SQLException e) {
            throw new RuntimeException("Failed to add embeddings", e);
        }
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
        List<String> ids = new ArrayList<>();
        for (int i = 0; i < embeddings.size(); i++) {
            ids.add(generateId());
        }
        addAll(ids, embeddings, textSegments);
        return ids;
    }

    @Override
    public List<EmbeddingMatch<TextSegment>> findRelevant(Embedding referenceEmbedding, int maxResults, double minScore) {
        try (Connection conn = getConnection()) {
            String sql = String.format(
                "SELECT id, text, metadata, l2_distance(embedding, ?) as distance " +
                "FROM %s " +
                "ORDER BY distance ASC " +
                "APPROXIMATE LIMIT ?",
                tableName
            );

            List<EmbeddingMatch<TextSegment>> results = new ArrayList<>();
            try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
                pstmt.setString(1, formatVector(referenceEmbedding.vector()));
                pstmt.setInt(2, maxResults);

                try (ResultSet rs = pstmt.executeQuery()) {
                    while (rs.next()) {
                        double score = 1.0 / (1.0 + rs.getDouble("distance")); // 将L2距离转换为相似度分数
                        if (score >= minScore) {
                            String id = rs.getString("id");
                            String text = rs.getString("text");
                            String metadataStr = rs.getString("metadata");

                            TextSegment segment = TextSegment.from(
                                text,
                                metadataStr != null ? parseMetadata(metadataStr) : null
                            );

                            results.add(new EmbeddingMatch<>(score, id, null, segment));
                        }
                    }
                }
            }
            return results;
        } catch (SQLException e) {
            throw new RuntimeException("Failed to find relevant embeddings", e);
        }
    }

    private Connection getConnection() throws SQLException {
        return DriverManager.getConnection(url, username, password);
    }

    private String generateId() {
        return java.util.UUID.randomUUID().toString();
    }

    @Override
    public void removeAll(java.util.Collection<String> ids) {
        if (ids == null || ids.isEmpty()) {
            throw new IllegalArgumentException("ids must not be empty");
        }
        try (Connection conn = getConnection()) {
            String sql = String.format(
                "DELETE FROM %s WHERE id IN (%s)",
                tableName,
                String.join(",", Collections.nCopies(ids.size(), "?"))
            );

            try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
                int i = 1;
                for (String id : ids) {
                    pstmt.setString(i++, id);
                }
                pstmt.executeUpdate();
            }
        } catch (SQLException e) {
            throw new RuntimeException("Failed to remove embeddings", e);
        }
    }

    @Override
    public void removeAll() {
        try (Connection conn = getConnection()) {
            try (Statement stmt = conn.createStatement()) {
                stmt.execute(String.format("TRUNCATE TABLE %s", tableName));
            }
        } catch (SQLException e) {
            throw new RuntimeException("Failed to remove all embeddings", e);
        }
    }

    private String formatVector(float[] vector) {
        StringBuilder sb = new StringBuilder();
        sb.append('[');
        for (int i = 0; i < vector.length; i++) {
            if (i > 0) sb.append(',');
            sb.append(vector[i]);
        }
        sb.append(']');
        return sb.toString();
    }

    private Metadata parseMetadata(String metadataStr) {
        Gson gson = new Gson();
        Map<String, String> metadataMap = gson.fromJson(metadataStr, new TypeToken<Map<String, String>>(){}.getType());
        return Metadata.from(metadataMap != null ? metadataMap : new HashMap<>());
    }
}